package com.akita.spring.ai.controller;

import com.akita.spring.ai.entity.vo.Result;
import com.akita.spring.ai.enums.BizTypeEnum;
import com.akita.spring.ai.repository.ChatHistoryRepository;
import com.akita.spring.ai.repository.FileRepository;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.BooleanUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.ParagraphPdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;

@Slf4j
@RestController
@RequestMapping("/ai/pdf")
public class PdfController {

    @Resource
    private FileRepository fileRepository;

    @Resource
    private VectorStore vectorStore;

    @Resource
    private ChatClient pdfChatClient;

    @Resource ChatClient aliyunChatClient;

    @Resource
    private ChatHistoryRepository chatHistoryRepository;

    @RequestMapping(value = "/chat", produces = "text/html;charset=UTF-8")
    public Flux<String> chat(String prompt, String chatId) {
        chatHistoryRepository.save(BizTypeEnum.PDF.getType(), chatId);

        org.springframework.core.io.Resource file = fileRepository.getFile(chatId);
        if (BooleanUtils.isFalse(file.exists())) {
            return Flux.just("文件不存在");
        }

        return pdfChatClient.prompt()
                .user(prompt)
                .advisors(a -> a.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .advisors(a -> a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION, "file_name == '" + file.getFilename() + "'"))
                .stream()
                .content();
    }

    /**
     * 上传PDF文件
     *
     * @param chatId 聊天ID
     * @param file   PDF文件
     * @return 上传结果
     */
    @RequestMapping("/upload/{chatId}")
    public Result uploadPdf(@PathVariable String chatId, @RequestParam MultipartFile file) {
        try {
            if (!Objects.equals(file.getContentType(), "application/pdf")) {
                return Result.fail("只能上传PDF文件");
            }
            Boolean success = fileRepository.save(chatId, file.getResource());
            if (BooleanUtils.isFalse(success)) {
                return Result.fail("文件上传失败");
            }
            this.writeToVectorStore(file.getResource());
            return Result.ok();
        } catch (Exception e) {
            log.error("文件上传失败", e);
            return Result.fail("文件上传失败");
        }
    }

    @GetMapping("/file/{chatId}")
    public ResponseEntity<org.springframework.core.io.Resource> download(@PathVariable("chatId") String chatId) {
        org.springframework.core.io.Resource resource = fileRepository.getFile(chatId);
        if (BooleanUtils.isFalse(resource.exists())) {
            return ResponseEntity.notFound().build();
        }
        String fileName = URLEncoder.encode(Objects.requireNonNull(resource.getFilename()), StandardCharsets.UTF_8);
        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_OCTET_STREAM)
                .header("Content-Disposition", "attachment; filename=\"" + fileName + "\"")
                .body(resource);
    }

    private void writeToVectorStore(org.springframework.core.io.Resource resource) {

//        ParagraphPdfDocumentReader paragraphPdfDocumentReader = new ParagraphPdfDocumentReader(resource);

        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                resource,
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1)
                        .build()
        );
        List<Document> documentList = reader.read();
        vectorStore.add(documentList);
    }

}
